import tensorflow as tf
from tensorflow.keras import layers, models

# 定义LeNet-5模型
def LeNet5():
    model = models.Sequential([
        # 第一层：卷积层
        layers.Conv2D(filters=6, kernel_size=(5, 5), activation='relu', input_shape=(32, 32, 1)),
        # 第二层：平均池化层
        layers.AveragePooling2D(pool_size=(2, 2)),
        # 第三层：卷积层
        layers.Conv2D(filters=16, kernel_size=(5, 5), activation='relu'),
        # 第四层：平均池化层
        layers.AveragePooling2D(pool_size=(2, 2)),
        # 将多维数据展平为一维
        layers.Flatten(),
        # 第五层：全连接层
        layers.Dense(120, activation='relu'),
        # 第六层：全连接层
        layers.Dense(84, activation='relu'),
        # 第七层：输出层，使用softmax激活函数进行多分类
        layers.Dense(10, activation='softmax')
    ])
    return model

if __name__ == "__main__":
    # 创建LeNet-5模型实例
    model = LeNet5()

    # 打印模型结构
    model.summary()

    # 编译模型
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
